from learning_to_adapt.utils.serializable import Serializable
from learning_to_adapt.utils.utils import remove_scope_from_name
from learning_to_adapt.dynamics.core.utils import *
import tensorflow as tf
import copy
from collections import OrderedDict

class Layer(Serializable):
    """
    一个容器，用于存储当前的更新前和更新后的策略
    同时提供执行和更新策略参数的功能

    注意：
        更新前的策略存储为tf.Variables，而更新后的策略存储在numpy数组中，并通过tf.placeholders执行

    参数：
        obs_dim (int)：观察空间的维度 -> 指定策略的输入大小
        action_dim (int)：动作空间的维度 -> 指定策略的输出大小
        name (str) ：用于给策略中的变量设置范围的名称
        hidden_sizes (tuple) ：网络隐藏层的大小
        learn_std (bool) ：是否学习网络输出的方差
        hidden_nonlinearity (Operation) ：隐藏层之间使用的非线性操作
        output_nonlinearity (Operation) ：最后一层网络之后使用的非线性操作
    """

    def __init__(self,
                 name,
                 input_dim,
                 output_dim,
                 hidden_sizes=(32, 32),
                 hidden_nonlinearity=tf.nn.relu,
                 output_nonlinearity=None,
                 input_var=None,
                 params=None,
                 **kwargs
                 ):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.name = name
        self.input_var = input_var
        self.hidden_sizes = hidden_sizes
        self.hidden_nonlinearity = hidden_nonlinearity
        self.output_nonlinearity = output_nonlinearity
        self.batch_normalization = kwargs.get('batch_normalization', False)
        self._params = params
        self._assign_ops = None
        self._assign_phs = None

    def build_graph(self):
        """
        构建策略的计算图
        """
        raise NotImplementedError

    """ --- 用于序列化的方法 --- """

    def get_params(self):
        """
        获取表示网络可训练权重（符号）的tf.Variables

        返回值：
            (dict)：包含所有可训练变量的字典
        """
        return self._params

    def get_param_values(self):
        """
        获取网络中所有当前权重的列表（在原始代码中进行了扁平化处理）

        返回值：
            (list)：参数的值列表
        """
        param_values = tf.get_default_session().run(self._params)
        return param_values

    def set_params(self, policy_params):
        """
        设置计算图的参数

        参数：
            policy_params (dict)：包含变量名称和相应参数值的字典
        """
        assert all([k1 == k2 for k1, k2 in zip(self.get_params().keys(), policy_params.keys())]), \
            "参数键必须与变量匹配"

        if self._assign_ops is None:
            assign_ops, assign_phs = [], []
            for var in self.get_params().values():
                assign_placeholder = tf.placeholder(dtype=var.dtype)
                assign_op = tf.assign(var, assign_placeholder)
                assign_ops.append(assign_op)
                assign_phs.append(assign_placeholder)
            self._assign_ops = assign_ops
            self._assign_phs = assign_phs
        feed_dict = dict(zip(self._assign_phs, policy_params.values()))
        tf.get_default_session().run(self._assign_ops, feed_dict=feed_dict)

    def __getstate__(self):
        state = {
            # 'init_args': Serializable.__getstate__(self),
            'network_params': self.get_param_values()
        }
        return state

    def __setstate__(self, state):
        # Serializable.__setstate__(self, state['init_args'])
        tf.get_default_session().run(tf.variables_initializer(self.get_params().values()))
        self.set_params(state['network_params'])

class MLP(Layer):
    """
    高斯多层感知器（MLP）策略（对角协方差矩阵）
    提供执行和更新策略参数的功能
    一个容器，用于存储当前的更新前和更新后的策略

    参数：
        obs_dim (int)：观察空间的维度 -> 指定策略的输入大小
        action_dim (int)：动作空间的维度 -> 指定策略的输出大小
        name (str)：策略名称，用作tf变量作用域
        hidden_sizes (tuple)：MLP的隐藏层大小的元组
        hidden_nonlinearity (tf.op)：隐藏层的非线性函数
        output_nonlinearity (tf.op或None)：输出层的非线性函数
        learn_std (布尔值)：标准差/方差是否可训练的变量还是固定的变量
        init_std (float)：初始策略标准差
        min_std (float)：最小策略标准差
    """

    def __init__(self, *args, **kwargs):
        # 存储序列化的初始参数，并调用超类的构造函数
        Serializable.quick_init(self, locals())
        Layer.__init__(self, *args, **kwargs)

        self.build_graph()

    def build_graph(self):
        """
        构建策略的计算图
        """
        with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):

            if self._params is None:
                # 构建实际的策略网络
                self.input_var, self.output_var = create_mlp(output_dim=self.output_dim,
                                                             hidden_sizes=self.hidden_sizes,
                                                             hidden_nonlinearity=self.hidden_nonlinearity,
                                                             output_nonlinearity=self.output_nonlinearity,
                                                             input_dim=(None, self.input_dim,),
                                                             input_var=self.input_var,
                                                             batch_normalization=self.batch_normalization,
                                                             )

                # 将可训练变量保存在字典中
                current_scope = tf.get_default_graph().get_name_scope()
                trainable_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=current_scope)
                self._params = OrderedDict([(remove_scope_from_name(var.name, current_scope), var)
                                            for var in trainable_vars])

            else:
                self.input_var, self.output_var = forward_mlp(output_dim=self.output_dim,
                                                              hidden_sizes=self.hidden_sizes,
                                                              hidden_nonlinearity=self.hidden_nonlinearity,
                                                              output_nonlinearity=self.output_nonlinearity,
                                                              input_var=self.input_var,
                                                              mlp_params=self._params)
class RNN(Layer):
    """
    高斯多层感知器（MLP）策略（对角协方差矩阵）
    提供执行和更新策略参数的功能
    一个容器，用于存储当前的更新前和更新后的策略

    参数：
        obs_dim (int)：观察空间的维度 -> 指定策略的输入大小
        action_dim (int)：动作空间的维度 -> 指定策略的输出大小
        name (str)：策略名称，用作tf变量作用域
        hidden_sizes (tuple)：MLP的隐藏层大小的元组
        hidden_nonlinearity (tf.op)：隐藏层的非线性函数
        output_nonlinearity (tf.op或None)：输出层的非线性函数
        learn_std (布尔值)：标准差/方差是否可训练的变量还是固定的变量
        init_std (float)：初始策略标准差
        min_std (float)：最小策略标准差
    """

    def __init__(self, *args, **kwargs):
        # 存储序列化的初始参数，并调用超类的构造函数
        self._cell_type = kwargs.get('cell_type', 'gru')
        self.state_var = kwargs.get('state_var', None)
        Layer.__init__(self, *args, **kwargs)

        self.build_graph()

    def build_graph(self):
        """
        构建策略的计算图
        """
        with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
            # 构建实际的策略网络
            args = create_rnn(name='rnn',
                              cell_type=self._cell_type,
                              output_dim=self.output_dim,
                              hidden_sizes=self.hidden_sizes,
                              hidden_nonlinearity=self.hidden_nonlinearity,
                              output_nonlinearity=self.output_nonlinearity,
                              input_dim=(None, None, self.input_dim,),
                              input_var=self.input_var,
                              state_var=self.state_var,
                              )

            self.input_var, self.state_var, self.output_var, self.next_state_var, self.cell = args

        current_scope = tf.get_default_graph().get_name_scope()
        trainable_policy_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=current_scope)
        self._params = OrderedDict([(remove_scope_from_name(var.name, current_scope), var) for var in trainable_policy_vars])


